from channels.generic.websocket import WebsocketConsumer
from django.http.request import QueryDict
from webssh.ssh import SSH
from asset.models import *
from utils.custom_log import log_start
logger = log_start('wss')
import json

# from django.utils.six import StringIO
# from io import StringIO

class WebSSH(WebsocketConsumer): 
    
    message = {'status': 0, 'message': None}
    """
    status:
        0: ssh 连接正常, websocket 正常
        1: 发生未知错误, 关闭 ssh 和 websocket 连接

    message:
        status 为 1 时, message 为具体的错误信息
        status 为 0 时, message 为 ssh 返回的数据, 前端页面将获取 ssh 返回的数据并写入终端页面
    """

    def connect(self):
        """
        打开 websocket 连接, 通过前端传入的参数尝试连接 ssh 主机
        :return:
        """
        try:
            self.accept()
            query_string = self.scope.get('query_string')
            ssh_args = QueryDict(query_string=query_string, encoding='utf-8')

            id = ssh_args.get('id') # 获取row id
            HostInstance = Host.objects.get(id=int(id))

            host = HostInstance.ip
            port = HostInstance.ssh_id.ssh_port
            user = HostInstance.ssh_id.ssh_user
            password = HostInstance.ssh_id.password
            private = HostInstance.ssh_id.private

            width = int(ssh_args.get('width'))
            height = int(ssh_args.get('height'))

            ssh_connect_dict = {
                'host': host,
                'user': user,
                'port': port,
                'password': password,
                'private':private,
                'width': width,
                'height': height
            }
            # logger.info(f'{ssh_connect_dict}')
            self.ssh = SSH(ssh_connect_dict, websocker=self, message=self.message)
        except Exception as e:
            message = {'status': 1, 'message': '获取主机信息失败,%s' % e}
            self.send(json.dumps(message))

    def disconnect(self, close_code):
        try:
            self.ssh.close_ssh()
        except Exception as  e:
            pass

    def receive(self, text_data=None, bytes_data=None):
        if text_data is None:
            self.ssh.django_to_ssh(bytes_data)
        else:
            data = json.loads(text_data)
            if type(data) == dict:
                status = data['status']
                if status == 0:
                    data = data['data']
                    print(data)
                    self.ssh.task_shell(data)
                else:
                    cols = data['cols']
                    rows = data['rows']
                    self.ssh.resize_pty(cols=cols, rows=rows)
